from build_gym import BipedalWheeledRobotEnv
from stable_baselines3 import PPO
from stable_baselines3.common.env_util import make_vec_env

if __name__ == "__main__":
    # 创建环境
    env = BipedalWheeledRobotEnv()

    # 使用 PPO 算法进行强化学习训练
    model = PPO("MlpPolicy", env, verbose=1)

    # 训练模型
    model.learn(total_timesteps=100000)

    # 保存模型
    model.save("ppo_bipedal_wheeled_robot")

    # 测试训练后的模型
    obs = env.reset()
    for i in range(1000):
        action, _states = model.predict(obs, deterministic=True)
        obs, reward, done, info = env.step(action)
        env.render()
        if done:
            obs = env.reset()

    env.close()
